{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Data loading pipeline examples\n",
    "\n",
    "The purpose of this notebook is to illustrate reading Nifti files and test speed of different methods."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "MONAI version: 0.1a1.dev8+6.gb3c5761.dirty\n",
      "Python version: 3.6.9 |Anaconda, Inc.| (default, Jul 30 2019, 19:07:31)  [GCC 7.3.0]\n",
      "Numpy version: 1.18.1\n",
      "Pytorch version: 1.4.0\n",
      "Ignite version: 0.3.0\n"
     ]
    }
   ],
   "source": [
    "%matplotlib inline\n",
    "\n",
    "import os\n",
    "import sys\n",
    "from glob import glob\n",
    "import tempfile\n",
    "\n",
    "import numpy as np\n",
    "import nibabel as nib\n",
    "\n",
    "\n",
    "import torch\n",
    "from torch.utils.data import DataLoader\n",
    "from torch.multiprocessing import Pool, Process, set_start_method\n",
    "try:\n",
    "     set_start_method('spawn')\n",
    "except RuntimeError:\n",
    "    pass\n",
    "\n",
    "import monai\n",
    "from monai.data import NiftiDataset\n",
    "from monai.transforms import (Compose, AddChannel, ScaleIntensity, ToTensor, \n",
    "                              RandSpatialCrop, Rotate, RandAffine)\n",
    "\n",
    "monai.config.print_config()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 0. Preparing input data (nifti images)\n",
    "\n",
    "Create a number of test Nifti files, 3d single channel images with spatial size (256, 256, 256) voxels."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "tempdir = tempfile.mkdtemp()\n",
    "\n",
    "for i in range(5):\n",
    "    im, seg = monai.data.synthetic.create_test_image_3d(256,256,256)\n",
    "    \n",
    "    n = nib.Nifti1Image(im, np.eye(4))\n",
    "    nib.save(n, os.path.join(tempdir, 'im%i.nii.gz'%i))\n",
    "    \n",
    "    n = nib.Nifti1Image(seg, np.eye(4))\n",
    "    nib.save(n, os.path.join(tempdir, 'seg%i.nii.gz'%i))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# prepare list of image names and segmentation names\n",
    "images = sorted(glob(os.path.join(tempdir,'im*.nii.gz')))\n",
    "segs = sorted(glob(os.path.join(tempdir,'seg*.nii.gz')))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1. Test image loading with minimal preprocessing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([3, 1, 256, 256, 256]) torch.Size([3, 1, 256, 256, 256])\n"
     ]
    }
   ],
   "source": [
    "imtrans = Compose([\n",
    "    AddChannel(),\n",
    "    ToTensor()\n",
    "])    \n",
    "\n",
    "segtrans = Compose([\n",
    "    AddChannel(),\n",
    "    ToTensor()\n",
    "])    \n",
    "    \n",
    "ds = NiftiDataset(images, segs, transform=imtrans, seg_transform=segtrans)\n",
    "loader = DataLoader(ds, batch_size=3, num_workers=8)\n",
    "\n",
    "im, seg = monai.utils.misc.first(loader)\n",
    "print(im.shape, seg.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "3.73 s ± 111 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)\n"
     ]
    }
   ],
   "source": [
    "%timeit data = next(iter(loader))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2. Test image-patch loading with CPU multi-processing:\n",
    "\n",
    "- rotate (256, 256, 256)-voxel in the plane axes=(1, 2)\n",
    "- extract random (64, 64, 64) patches\n",
    "- implemented in MONAI using ` scipy.ndimage.rotate`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([3, 1, 64, 64, 64]) torch.Size([3, 1, 64, 64, 64])\n"
     ]
    }
   ],
   "source": [
    "images = sorted(glob(os.path.join(tempdir,'im*.nii.gz')))\n",
    "segs = sorted(glob(os.path.join(tempdir,'seg*.nii.gz')))\n",
    "\n",
    "imtrans = Compose([\n",
    "    ScaleIntensity(),\n",
    "    AddChannel(),\n",
    "    Rotate(angle=45.),\n",
    "    RandSpatialCrop((64, 64, 64), random_size=False),\n",
    "    ToTensor()\n",
    "])    \n",
    "\n",
    "segtrans = Compose([\n",
    "    AddChannel(),\n",
    "    Rotate(angle=45.),\n",
    "    RandSpatialCrop((64, 64, 64), random_size=False),\n",
    "    ToTensor()\n",
    "])    \n",
    "    \n",
    "ds = NiftiDataset(images, segs, transform=imtrans, seg_transform=segtrans)\n",
    "loader = DataLoader(ds, batch_size=3, num_workers=8, pin_memory=torch.cuda.is_available())\n",
    "\n",
    "im, seg = monai.utils.misc.first(loader)\n",
    "print(im.shape, seg.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "10.1 s ± 83.1 ms per loop (mean ± std. dev. of 7 runs, 3 loops each)\n"
     ]
    }
   ],
   "source": [
    "%timeit -n 3 data = next(iter(loader))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "(the above results were based on a 2.9 GHz 6-Core Intel Core i9)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3. Test image-patch loading with preprocessing on GPU:\n",
    "\n",
    "- random rotate (256, 256, 256)-voxel in the plane axes=(1, 2)\n",
    "- extract random (64, 64, 64) patches\n",
    "- implemented in MONAI using native pytorch resampling"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([3, 1, 64, 64, 64]) torch.Size([3, 1, 64, 64, 64])\n"
     ]
    }
   ],
   "source": [
    "images = sorted(glob(os.path.join(tempdir,'im*.nii.gz')))\n",
    "segs = sorted(glob(os.path.join(tempdir,'seg*.nii.gz')))\n",
    "\n",
    "# same parameter with different interpolation mode for image and segmentation\n",
    "rand_affine_img = RandAffine(prob=1.0, rotate_range=np.pi/4, translate_range=(96, 96, 96),\n",
    "                             spatial_size=(64, 64, 64), mode='bilinear',\n",
    "                             as_tensor_output=True, device=torch.device('cuda:0'))\n",
    "rand_affine_seg = RandAffine(prob=1.0, rotate_range=np.pi/4, translate_range=(96, 96, 96),\n",
    "                             spatial_size=(64, 64, 64), mode='nearest',\n",
    "                             as_tensor_output=True, device=torch.device('cuda:0'))\n",
    "    \n",
    "imtrans = Compose([\n",
    "    ScaleIntensity(),\n",
    "    AddChannel(),\n",
    "    rand_affine_img,\n",
    "    ToTensor()\n",
    "])    \n",
    "\n",
    "segtrans = Compose([\n",
    "    AddChannel(),\n",
    "    rand_affine_seg,\n",
    "    ToTensor()\n",
    "])    \n",
    "    \n",
    "ds = NiftiDataset(images, segs, transform=imtrans, seg_transform=segtrans)\n",
    "loader = DataLoader(ds, batch_size=3, num_workers=0)\n",
    "\n",
    "im, seg = monai.utils.misc.first(loader)\n",
    "\n",
    "print(im.shape, seg.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2.05 s ± 668 µs per loop (mean ± std. dev. of 7 runs, 3 loops each)\n"
     ]
    }
   ],
   "source": [
    "%timeit -n 3 data = next(iter(loader))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "TITAN Xp COLLECTORS EDITION\n",
      "|===========================================================================|\n",
      "|                  PyTorch CUDA memory summary, device ID 0                 |\n",
      "|---------------------------------------------------------------------------|\n",
      "|            CUDA OOMs: 0            |        cudaMalloc retries: 0         |\n",
      "|===========================================================================|\n",
      "|        Metric         | Cur Usage  | Peak Usage | Tot Alloc  | Tot Freed  |\n",
      "|---------------------------------------------------------------------------|\n",
      "| Allocated memory      |    6144 KB |  157696 KB |   17250 MB |   17244 MB |\n",
      "|---------------------------------------------------------------------------|\n",
      "| Active memory         |    6144 KB |  157696 KB |   17250 MB |   17244 MB |\n",
      "|---------------------------------------------------------------------------|\n",
      "| GPU reserved memory   |  225280 KB |  225280 KB |  225280 KB |       0 B  |\n",
      "|---------------------------------------------------------------------------|\n",
      "| Non-releasable memory |   14336 KB |   77823 KB |   11789 MB |   11775 MB |\n",
      "|---------------------------------------------------------------------------|\n",
      "| Allocations           |       2    |      14    |    2354    |    2352    |\n",
      "|---------------------------------------------------------------------------|\n",
      "| Active allocs         |       2    |      14    |    2354    |    2352    |\n",
      "|---------------------------------------------------------------------------|\n",
      "| GPU reserved segments |       8    |       8    |       8    |       0    |\n",
      "|---------------------------------------------------------------------------|\n",
      "| Non-releasable allocs |       1    |       6    |    1508    |    1507    |\n",
      "|===========================================================================|\n",
      "\n"
     ]
    }
   ],
   "source": [
    "print(torch.cuda.get_device_name(0))\n",
    "print(torch.cuda.memory_summary(0, abbreviated=True))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "!rm -rf {tempdir}"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}